import argparse
import textwrap
import binascii
import sys
import os
import shutil

from tools import *

def get_args():
    parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter,
                                     description=textwrap.dedent('''
                                     A tool to pack image with new structure'''))
    parser.add_argument('-raw_img', required=False, dest='raw', help='INPUT: The raw image')
    parser.add_argument('-out_img', required=False, dest='out',
                        help='OUTPUT: The processed image. Filename can be same to input image')
    parser.add_argument('-platform', required=False, dest='platform', choices=['hi1910'],
                        help='INPUT: platform : hi1910')

    # input for soc
    parser.add_argument('-root', required=False, dest='root', help='INPUT: CA cert, Parse N/E from it')
    parser.add_argument('--root_raw', help='root key whether is pem file', action="store_true")
    parser.add_argument('-cert', required=False, dest='cert', help='INPUT: The soc cert')
    parser.add_argument('-sig', required=False, dest='sig', help='INPUT: The soc signature')

    # input for bcm
    parser.add_argument('-bcmcert', required=False, dest='bcmcert', help='INPUT: The bcm cert')

    # input for cms
    parser.add_argument('-cms', required=False, dest='cms', help='INPUT: The cms file')
    parser.add_argument('-ini', required=False, dest='ini', help='INPUT: The ini file')
    parser.add_argument('-crl', required=False, dest='crl', help='INPUT: The crl file')

    # 1910_version
    parser.add_argument('-version', required=False, dest='ver', help='INPUT: The version number')

    # flag cmd
    parser.add_argument('--addsoc', help="choose whether add soc", action="store_true")
    parser.add_argument('-S', help="choose whether Onchiprom", action="store_true")
    
    parser.add_argument('-D', help="choose whether double ", action="store_true")
    parser.add_argument('-B', help="choose whether bcm", action="store_true")

    parser.add_argument('--addcms', help="choose whether add cms", action="store_true")

    # encrypt related
    parser.add_argument('-enc', choices=['aes', 'sm4'], help="enc mode")
    parser.add_argument('-key', required=False, dest='key', help='INPUT: The encrypt key')
    parser.add_argument('--pss', help="choose the code signature algo", action="store_true")

    # TA related
    parser.add_argument('--ta', required=False, dest='ta', help='INPUT: The name of optee signed TA, like xxxx.ta')
    parser.add_argument('--ta_path', required=False, dest='ta_path', help='INPUT: The path of TA and signed TA ELF')

    return parser.parse_args()


def main():
    if args.ta:
        sh_ta = os.path.join(args.ta_path, args.ta)
        st_ta = os.path.join(args.ta_path, args.ta[:-3] + '.stripped.elf')
        tmp_file = os.path.join(args.ta_path, args.ta + '.tmp')
        with open(tmp_file, "wb+") as o_f:
            with open(sh_ta, "rb") as ta_h:
                with open(st_ta, "rb") as ta_t:
                    ta = ta_h.read()
                    o_f.write(ta)
                    ta_t.seek(-552, 2); # siglen + sizeof(struct module_signature) + sizeof(MODULE_SIG_STRING)
                    sig = ta_t.read(512) # RSA4096 sig len
                    o_f.write(sig)
                    tl = construct_ta_tailer()
                    o_f.write(tl)

        shutil.copyfile(tmp_file, sh_ta)
    else:
        with open(args.raw, "rb") as f:
            hash_buf = cal_image_hash(f)
            code_len = get_filelen(f)

            tmp_file = args.out + '.tmp'
            with open(tmp_file, 'wb+') as o_f:
                platforms.write_header_huawei(args, o_f, hash_buf, code_len)
                if args.D:
                    platforms.write_header_customer(args, o_f, hash_buf, code_len)
                elif args.B:
                    print(args.bcmcert)
                    platforms.write_header_bcm(args.bcmcert, o_f)
                platforms.write_image(args, o_f)
                platforms.write_cms(args, o_f, code_len)
                platforms.write_extern(args, o_f, [hash_buf, code_len])

        shutil.copyfile(tmp_file, args.out)
    if os.path.exists(tmp_file):
        os.remove(tmp_file)


if __name__ == '__main__':
    args = get_args()
    if args.platform == 'hi1910':
        import hi_platform.platform_hi1910 as platforms
    elif args.ta == '':
        raise ValueError("error platform type {0}".format(args.platform))
    main()
